library(tidyverse)
library(maptools)
library(igraph)
library(parallel)
library(redist)
library(spdep)

ipw <- function(x, beta, pop){
    indpop <- which(x$distance_parity <= pop)
    indbeta <- which(x$beta_sequence == beta)
    inds <- intersect(indpop, indbeta)
    psi <- x$constraint_pop[inds]
    w <- 1 / exp(beta * psi)
    inds <- sample(inds, length(inds), replace = TRUE, prob = w)
    x <- x$partitions[,inds]
    return(x)
}
ben_theme <- function(){
    theme_classic() +
        theme(panel.background = element_blank(),
              panel.grid.major = element_blank(),
              panel.grid.minor = element_blank(),
              axis.line = element_line(colour = "black"),
              panel.border = element_rect(colour = "black", fill = NA, size = 1),
              strip.background = element_blank(),
              legend.position = "bottom", legend.title = element_blank(),
              plot.title = element_text(hjust = 0.5),
              plot.subtitle = element_text(hjust = .5))
}

fl_map <- readShapePoly("../../data/fl/FL.shp")

## Load nodes
nodes <- strsplit(readLines("../../data/largemap_enum_sample/map_sub_nodes.dat")," ")[[1]]
## nodes <- strsplit(readLines("../../data/big_validation_map/fl_sub_nodes.dat")," ")[[1]]

## First, subset down to right map
fl_sub <- fl_map[which(rownames(fl_map@data) %in% nodes),]

## Then, reorder nodes to be correct order to match - otherwise, won't
## line up properly with original map
fl_sub <- fl_sub[match(nodes, rownames(fl_sub@data)),]

## Convert factor columns to characters, because this particular data is messy
indx <- sapply(fl_sub@data, is.factor)
fl_sub@data[indx] <- lapply(
    fl_sub@data[indx], function(x) as.numeric(as.character(x))
)

## Get info on job
i <- as.numeric(Sys.getenv("SLURM_ARRAY_TASK_ID")) 

## --------------
## Run redist.rsg
## --------------
adjlist <- poly2nb(fl_sub, queen = FALSE)
adjlist <- lapply(adjlist, function(x){x-1})
nsims <- 10000

## Unconstrained
sv_out <- lapply(1:nsims, function(x){
    if(x %% (nsims / 10) == 0){
        print(paste0("Done with ", x, " no parity iterations at ", Sys.time(), "."))
    }
    return(redist.rsg(
        adjlist, fl_sub@data$pop, ndists = 2, thresh = 100, verbose = FALSE,
        maxiter = 500000
    )$district_membership)
})
rsg_out <- do.call(cbind, sv_out)
rsg_seg_full <- redist.segcalc(rsg_out, grouppop = fl_sub@data$mccain, 
                               fullpop = fl_sub@data$pop)

## 5%
sv_out <- lapply(1:nsims, function(x){
    if(x %% (nsims / 10) == 0){
        print(paste0("Done with ", x, " 5% iterations at ", Sys.time(), "."))
    }
    return(redist.rsg(
        adjlist, fl_sub@data$pop, ndists = 2, thresh = .05, verbose = FALSE,
        maxiter = 500000
    )$district_membership)
})
rsg_out <- do.call(cbind, sv_out)
rsg_seg_05p <- redist.segcalc(rsg_out, grouppop = fl_sub@data$mccain, 
                              fullpop = fl_sub@data$pop)

## 1%
sv_out <- lapply(1:nsims, function(x){
    if(x %% (nsims / 10) == 0){
        print(paste0("Done with ", x, " 1% iterations at ", Sys.time(), "."))
    }
    return(redist.rsg(
        adjlist, fl_sub@data$pop, ndists = 2, thresh = .01, verbose = FALSE,
        maxiter = 500000
    )$district_membership)
})
rsg_out <- do.call(cbind, sv_out)
rsg_seg_01p <- redist.segcalc(rsg_out, grouppop = fl_sub@data$mccain, 
                              fullpop = fl_sub@data$pop)

## -------------
## Create output
## -------------
df_out <- data.frame(seg = c(rsg_seg_full, rsg_seg_05p, rsg_seg_01p),
                     parity = c(rep("No Equal Population Constraint", nsims),
                                rep("5% Equal Population Constraint", nsims),
                                rep("1% Equal Population Constraint", nsims)))
write_csv(df_out, path = paste0("../../data/largemap_enum_sample/rsg_", i, ".csv"))
